#include <x86.h>
#include <string.h>
#define PG_ALIGN __attribute((aligned(PGSIZE)))

static PDE kpdirs[NR_PDE] PG_ALIGN;
static PTE kptabs[PMEM_SIZE / PGSIZE] PG_ALIGN;
static void* (*palloc_f)();
static void (*pfree_f)(void*);

_Area segments[] = {      // Kernel memory mappings
  {.start = (void*)0,          .end = (void*)PMEM_SIZE}
};

#define NR_KSEG_MAP (sizeof(segments) / sizeof(segments[0]))

void _pte_init(void* (*palloc)(), void (*pfree)(void*)) {
  palloc_f = palloc;
  pfree_f = pfree;

  int i;

  // make all PDEs invalid
  for (i = 0; i < NR_PDE; i ++) {
    kpdirs[i] = 0;
  }

  PTE *ptab = kptabs;
  for (i = 0; i < NR_KSEG_MAP; i ++) {
    uint32_t pdir_idx = (uintptr_t)segments[i].start / (PGSIZE * NR_PTE);
    uint32_t pdir_idx_end = (uintptr_t)segments[i].end / (PGSIZE * NR_PTE);
    for (; pdir_idx < pdir_idx_end; pdir_idx ++) {
      // fill PDE
      kpdirs[pdir_idx] = (uintptr_t)ptab | PTE_P;

      // fill PTE
      PTE pte = PGADDR(pdir_idx, 0, 0) | PTE_P;
      PTE pte_end = PGADDR(pdir_idx + 1, 0, 0) | PTE_P;
      for (; pte < pte_end; pte += PGSIZE) {
        *ptab = pte;
        ptab ++;
      }
    }
  }

  set_cr3(kpdirs);
  set_cr0(get_cr0() | CR0_PG);
}

void _protect(_Protect *p) {
  PDE *updir = (PDE*)(palloc_f());
  p->ptr = updir;
  // map kernel space
  for (int i = 0; i < NR_PDE; i ++) {
    updir[i] = kpdirs[i];
  }

  p->area.start = (void*)0x8000000;
  p->area.end = (void*)0xc0000000;
}

void _release(_Protect *p) {
}

void _switch(_Protect *p) {
  set_cr3(p->ptr);
}

void _map(_Protect *p, void *va, void *pa) {
  // 获取页目录的基地址
    PDE *pdir = (PDE *)p->ptr;
    
    // 计算页目录索引和页表索引
    uint32_t pde_idx = PDX(va);
    uint32_t pte_idx = PTX(va);
    
    // 检查页目录项是否有效
    if (!(pdir[pde_idx] & PTE_P)) {
        // 页目录项无效，分配新的页表
        PTE *ptab = (PTE *)palloc_f();
        pdir[pde_idx] = (uint32_t)ptab | PTE_P;
    }
    
    // 获取页表基地址
    PTE *ptab = (PTE *)PTE_ADDR(pdir[pde_idx]);
    
    // 填写页表项
    ptab[pte_idx] = (uint32_t)pa | PTE_P | PTE_W | PTE_U;
}

void _unmap(_Protect *p, void *va) {
}

_RegSet *_umake(_Protect *p, _Area ustack, _Area kstack, void *entry, char *const argv[], char *const envp[]) {
  // 用户栈从高地址向低地址增长，栈底是 ustack.end
  uintptr_t stack_top = (uintptr_t)ustack.end;

  // 1. 在栈底预留陷阱帧（_RegSet）的空间
  _RegSet *tf = (_RegSet*)(stack_top - sizeof(_RegSet));
  
  // 2. 在陷阱帧上方预留 _start() 的参数栈帧（argc, argv, envp）
  uintptr_t *args_frame = (uintptr_t*)((uintptr_t)tf - 3 * sizeof(uintptr_t));
  
  // 初始化参数栈帧：argc=0, argv=NULL, envp=NULL
  args_frame[0] = 0;     // argc
  args_frame[1] = 0;     // argv (NULL)
  args_frame[2] = 0;     // envp (NULL)

  // 3. 初始化陷阱帧：
  memset(tf, 0, sizeof(_RegSet)); // 所有寄存器初始化为0
  tf->eflags =  FL_IF;
  tf->cs = 8;
  tf->eip = (uintptr_t)entry;     // 返回地址设置为 entry
  tf->esp = (uintptr_t)args_frame;// 栈指针指向参数栈帧的起始位置

  return tf; // 返回陷阱帧指针，由 Nanos-lite 保存到 PCB
}
